package com.wechat.gateway.ratelimit;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.util.StreamUtils;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import jakarta.annotation.PostConstruct;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 基于Redis Lua脚本的令牌桶管理器
 */
@Slf4j
@Component
public class TokenBucketManager {

    private static final String BUCKET_KEY_PREFIX = "rate_limit:bucket:";
    private static final long DEFAULT_EXPIRE_SECONDS = 3600; // 1小时过期

    // Lua脚本文件路径
    private static final String LUA_SCRIPTS_PATH = "lua/";
    private static final String CONSUME_TOKENS_SCRIPT = "consume-tokens.lua";
    private static final String GET_BUCKET_SCRIPT = "get-bucket.lua";
    private static final String RESET_BUCKET_SCRIPT = "reset-bucket.lua";
    private static final String DELETE_BUCKET_SCRIPT = "delete-bucket.lua";

    @Autowired
    private StringRedisTemplate stringRedisTemplate;
    @Autowired
    private ObjectMapper objectMapper;

    // Redis脚本
    private DefaultRedisScript<List> consumeTokensScript;
    private DefaultRedisScript<String> getBucketScript;
    private DefaultRedisScript<Long> resetBucketScript;
    private DefaultRedisScript<Long> deleteBucketScript;

    @PostConstruct
    public void initScripts() {
        try {
            // 从文件加载令牌消费脚本
            consumeTokensScript = new DefaultRedisScript<>();
            consumeTokensScript.setScriptText(loadLuaScript(CONSUME_TOKENS_SCRIPT));
            consumeTokensScript.setResultType(List.class);

            // 从文件加载获取桶状态脚本
            getBucketScript = new DefaultRedisScript<>();
            getBucketScript.setScriptText(loadLuaScript(GET_BUCKET_SCRIPT));
            getBucketScript.setResultType(String.class);

            // 从文件加载重置桶脚本
            resetBucketScript = new DefaultRedisScript<>();
            resetBucketScript.setScriptText(loadLuaScript(RESET_BUCKET_SCRIPT));
            resetBucketScript.setResultType(Long.class);

            // 从文件加载删除桶脚本
            deleteBucketScript = new DefaultRedisScript<>();
            deleteBucketScript.setScriptText(loadLuaScript(DELETE_BUCKET_SCRIPT));
            deleteBucketScript.setResultType(Long.class);

            log.info("Redis Lua脚本从文件加载完成: {}",
                    List.of(CONSUME_TOKENS_SCRIPT, GET_BUCKET_SCRIPT, RESET_BUCKET_SCRIPT, DELETE_BUCKET_SCRIPT));
        } catch (Exception e) {
            log.error("Redis Lua脚本初始化失败", e);
            throw new RuntimeException("Lua脚本加载失败", e);
        }
    }

    /**
     * 从classpath加载Lua脚本文件
     */
    private String loadLuaScript(String scriptName) throws IOException {
        String scriptPath = LUA_SCRIPTS_PATH + scriptName;
        ClassPathResource resource = new ClassPathResource(scriptPath);

        if (!resource.exists()) {
            throw new IOException("Lua脚本文件不存在: " + scriptPath);
        }

        String script = StreamUtils.copyToString(resource.getInputStream(), StandardCharsets.UTF_8);
        log.debug("加载Lua脚本: {} ({}字符)", scriptName, script.length());
        return script;
    }

    /**
     * 尝试消费令牌 - 使用Redis Lua脚本和JSON存储
     */
    public Mono<TokenBucketResult> tryConsume(String bucketKey, RateLimitConfig config, long requestedTokens) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;
        long now = Instant.now().getEpochSecond();

        return Mono.fromCallable(() -> {
            try {
                // 执行Lua脚本：capacity, refillRate, requestedTokens, now, expireSeconds
                List<String> scriptArgs = Arrays.asList(
                        String.valueOf(config.getCapacity()),
                        String.valueOf(config.getRefillRate()),
                        String.valueOf(requestedTokens),
                        String.valueOf(now),
                        String.valueOf(DEFAULT_EXPIRE_SECONDS));

                @SuppressWarnings("unchecked")
                List<Long> result = stringRedisTemplate.execute(
                        consumeTokensScript,
                        List.of(redisKey),
                        scriptArgs.toArray());

                if (result == null || result.size() < 3) {
                    log.warn("Redis脚本返回结果异常: bucketKey={}, result={}", bucketKey, result);
                    return createFailsafeResult(bucketKey, config.getRuleName());
                }

                boolean allowed = result.get(0) == 1L;
                long remainingTokens = result.get(1);
                long waitTime = result.get(2);

                if (log.isDebugEnabled()) {
                    log.debug("令牌桶消费结果: rule={}, bucketKey={}, allowed={}, remaining={}, waitTime={}",
                            config.getRuleName(), bucketKey, allowed, remainingTokens, waitTime);
                }

                return TokenBucketResult.builder()
                        .allowed(allowed)
                        .remainingTokens(remainingTokens)
                        .waitTime(waitTime)
                        .ruleName(config.getRuleName())
                        .bucketKey(bucketKey)
                        .build();
            } catch (Exception e) {
                log.error("令牌桶操作失败: bucketKey={}, rule={}, error={}", bucketKey, config.getRuleName(), e.getMessage(),
                        e);
                return createFailsafeResult(bucketKey, config.getRuleName());
            }
        })
                .subscribeOn(Schedulers.boundedElastic());
    }

    /**
     * 获取桶状态 - 返回JSON格式的桶信息
     */
    public Mono<TokenBucket> getBucketStatus(String bucketKey) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;
        long now = Instant.now().getEpochSecond();

        return Mono.fromCallable(() -> {
            try {
                String result = stringRedisTemplate.execute(
                        getBucketScript,
                        Arrays.asList(redisKey),
                        String.valueOf(now));

                if (result == null || result.isEmpty()) {
                    log.debug("桶不存在: {}", bucketKey);
                    return null;
                }

                return parseTokenBucketFromJson(result);
            } catch (Exception e) {
                log.error("获取桶状态失败: bucketKey={}, error={}", bucketKey, e.getMessage(), e);
                return null;
            }
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(null);
    }

    /**
     * 重置桶状态
     */
    public Mono<Boolean> resetBucket(String bucketKey, RateLimitConfig config) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;
        long now = Instant.now().getEpochSecond();

        return Mono.fromCallable(() -> {
            try {
                Long result = stringRedisTemplate.execute(
                        resetBucketScript,
                        Arrays.asList(redisKey),
                        String.valueOf(config.getCapacity()),
                        String.valueOf(config.getRefillRate()),
                        String.valueOf(now),
                        String.valueOf(DEFAULT_EXPIRE_SECONDS));

                boolean success = result != null && result == 1L;
                if (success) {
                    log.info("重置令牌桶成功: bucketKey={}, capacity={}, refillRate={}",
                            bucketKey, config.getCapacity(), config.getRefillRate());
                } else {
                    log.warn("重置令牌桶失败: bucketKey={}", bucketKey);
                }
                return success;
            } catch (Exception e) {
                log.error("重置令牌桶失败: bucketKey={}, error={}", bucketKey, e.getMessage(), e);
                return false;
            }
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(false);
    }

    /**
     * 删除桶
     */
    public Mono<Boolean> deleteBucket(String bucketKey) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;

        return Mono.fromCallable(() -> {
            try {
                Long result = stringRedisTemplate.execute(
                        deleteBucketScript,
                        Arrays.asList(redisKey));

                boolean deleted = result != null && result == 1L;
                if (deleted) {
                    log.info("删除令牌桶成功: {}", bucketKey);
                } else {
                    log.debug("令牌桶不存在或已删除: {}", bucketKey);
                }
                return deleted;
            } catch (Exception e) {
                log.error("删除令牌桶失败: bucketKey={}, error={}", bucketKey, e.getMessage(), e);
                return false;
            }
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(false);
    }

    /**
     * 检查桶是否存在
     */
    public Mono<Boolean> bucketExists(String bucketKey) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;
        return Mono.fromCallable(() -> {
            Boolean exists = stringRedisTemplate.hasKey(redisKey);
            log.debug("检查桶是否存在: bucketKey={}, exists={}", bucketKey, exists);
            return exists != null && exists;
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(false);
    }

    /**
     * 设置桶的TTL
     */
    public Mono<Boolean> setBucketTTL(String bucketKey, long seconds) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;
        return Mono.fromCallable(() -> {
            Boolean result = stringRedisTemplate.expire(redisKey, seconds, TimeUnit.SECONDS);
            log.debug("设置桶TTL: bucketKey={}, seconds={}, result={}", bucketKey, seconds, result);
            return result != null && result;
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(false);
    }

    /**
     * 获取桶的剩余TTL
     */
    public Mono<Long> getBucketTTL(String bucketKey) {
        String redisKey = BUCKET_KEY_PREFIX + bucketKey;
        return Mono.fromCallable(() -> {
            Long ttl = stringRedisTemplate.getExpire(redisKey, TimeUnit.SECONDS);
            log.debug("获取桶TTL: bucketKey={}, ttl={}", bucketKey, ttl);
            return ttl != null ? ttl : -1L;
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(-1L);
    }

    /**
     * 批量获取桶状态
     */
    public Mono<List<TokenBucket>> getBatchBucketStatus(List<String> bucketKeys) {
        return Mono.fromCallable(() -> {
            return bucketKeys.stream()
                    .map(key -> getBucketStatus(key).block())
                    .filter(bucket -> bucket != null)
                    .toList();
        })
                .subscribeOn(Schedulers.boundedElastic())
                .onErrorReturn(List.of());
    }

    // ==================== 私有方法 ====================

    /**
     * 创建故障安全结果（允许通过）
     */
    private TokenBucketResult createFailsafeResult(String bucketKey, String ruleName) {
        log.warn("使用故障安全模式: bucketKey={}, ruleName={}", bucketKey, ruleName);
        return TokenBucketResult.builder()
                .allowed(true) // 故障时允许通过
                .remainingTokens(0)
                .waitTime(0)
                .ruleName(ruleName)
                .bucketKey(bucketKey)
                .build();
    }

    /**
     * 从JSON字符串解析TokenBucket对象
     */
    private TokenBucket parseTokenBucketFromJson(String jsonData) {
        if (jsonData == null || jsonData.trim().isEmpty()) {
            return null;
        }
        try {
            TokenBucket bucket = objectMapper.readValue(jsonData, TokenBucket.class);
            log.debug("解析令牌桶JSON成功: capacity={}, tokens={}, refillRate={}",
                    bucket.getCapacity(), bucket.getTokens(), bucket.getRefillRate());
            return bucket;
        } catch (JsonProcessingException e) {
            log.error("解析令牌桶JSON失败: data={}, error={}", jsonData, e.getMessage());
            return null;
        }
    }

    /**
     * 将TokenBucket对象序列化为JSON
     */
    private String serializeTokenBucketToJson(TokenBucket bucket) {
        try {
            return objectMapper.writeValueAsString(bucket);
        } catch (JsonProcessingException e) {
            log.error("序列化令牌桶JSON失败: bucket={}, error={}", bucket, e.getMessage());
            throw new RuntimeException("TokenBucket序列化失败", e);
        }
    }
}